import torch
import numpy as np
from typing import List
from easyvolcap.utils.console_utils import *
from easyvolcap.utils.data_utils import to_cuda

# /home/xuzhen/miniconda3/envs/easyvv/lib/python3.10/site-packages/matplotlib/_cm_listed.py

jet = {'red': ((0.00, 0, 0),
               (0.35, 0, 0),
               (0.66, 1, 1),
               (0.89, 1, 1),
               (1.00, 0.5, 0.5)),
       'green': ((0.000, 0, 0),
                 (0.125, 0, 0),
                 (0.375, 1, 1),
                 (0.640, 1, 1),
                 (0.910, 0, 0),
                 (1.000, 0, 0)),
       'blue': ((0.00, 0.5, 0.5),
                (0.11, 1, 1),
                (0.34, 1, 1),
                (0.65, 0, 0),
                (1.00, 0, 0))}


virdis = [[0.267004, 0.004874, 0.329415],
          [0.268510, 0.009605, 0.335427],
          [0.269944, 0.014625, 0.341379],
          [0.271305, 0.019942, 0.347269],
          [0.272594, 0.025563, 0.353093],
          [0.273809, 0.031497, 0.358853],
          [0.274952, 0.037752, 0.364543],
          [0.276022, 0.044167, 0.370164],
          [0.277018, 0.050344, 0.375715],
          [0.277941, 0.056324, 0.381191],
          [0.278791, 0.062145, 0.386592],
          [0.279566, 0.067836, 0.391917],
          [0.280267, 0.073417, 0.397163],
          [0.280894, 0.078907, 0.402329],
          [0.281446, 0.084320, 0.407414],
          [0.281924, 0.089666, 0.412415],
          [0.282327, 0.094955, 0.417331],
          [0.282656, 0.100196, 0.422160],
          [0.282910, 0.105393, 0.426902],
          [0.283091, 0.110553, 0.431554],
          [0.283197, 0.115680, 0.436115],
          [0.283229, 0.120777, 0.440584],
          [0.283187, 0.125848, 0.444960],
          [0.283072, 0.130895, 0.449241],
          [0.282884, 0.135920, 0.453427],
          [0.282623, 0.140926, 0.457517],
          [0.282290, 0.145912, 0.461510],
          [0.281887, 0.150881, 0.465405],
          [0.281412, 0.155834, 0.469201],
          [0.280868, 0.160771, 0.472899],
          [0.280255, 0.165693, 0.476498],
          [0.279574, 0.170599, 0.479997],
          [0.278826, 0.175490, 0.483397],
          [0.278012, 0.180367, 0.486697],
          [0.277134, 0.185228, 0.489898],
          [0.276194, 0.190074, 0.493001],
          [0.275191, 0.194905, 0.496005],
          [0.274128, 0.199721, 0.498911],
          [0.273006, 0.204520, 0.501721],
          [0.271828, 0.209303, 0.504434],
          [0.270595, 0.214069, 0.507052],
          [0.269308, 0.218818, 0.509577],
          [0.267968, 0.223549, 0.512008],
          [0.266580, 0.228262, 0.514349],
          [0.265145, 0.232956, 0.516599],
          [0.263663, 0.237631, 0.518762],
          [0.262138, 0.242286, 0.520837],
          [0.260571, 0.246922, 0.522828],
          [0.258965, 0.251537, 0.524736],
          [0.257322, 0.256130, 0.526563],
          [0.255645, 0.260703, 0.528312],
          [0.253935, 0.265254, 0.529983],
          [0.252194, 0.269783, 0.531579],
          [0.250425, 0.274290, 0.533103],
          [0.248629, 0.278775, 0.534556],
          [0.246811, 0.283237, 0.535941],
          [0.244972, 0.287675, 0.537260],
          [0.243113, 0.292092, 0.538516],
          [0.241237, 0.296485, 0.539709],
          [0.239346, 0.300855, 0.540844],
          [0.237441, 0.305202, 0.541921],
          [0.235526, 0.309527, 0.542944],
          [0.233603, 0.313828, 0.543914],
          [0.231674, 0.318106, 0.544834],
          [0.229739, 0.322361, 0.545706],
          [0.227802, 0.326594, 0.546532],
          [0.225863, 0.330805, 0.547314],
          [0.223925, 0.334994, 0.548053],
          [0.221989, 0.339161, 0.548752],
          [0.220057, 0.343307, 0.549413],
          [0.218130, 0.347432, 0.550038],
          [0.216210, 0.351535, 0.550627],
          [0.214298, 0.355619, 0.551184],
          [0.212395, 0.359683, 0.551710],
          [0.210503, 0.363727, 0.552206],
          [0.208623, 0.367752, 0.552675],
          [0.206756, 0.371758, 0.553117],
          [0.204903, 0.375746, 0.553533],
          [0.203063, 0.379716, 0.553925],
          [0.201239, 0.383670, 0.554294],
          [0.199430, 0.387607, 0.554642],
          [0.197636, 0.391528, 0.554969],
          [0.195860, 0.395433, 0.555276],
          [0.194100, 0.399323, 0.555565],
          [0.192357, 0.403199, 0.555836],
          [0.190631, 0.407061, 0.556089],
          [0.188923, 0.410910, 0.556326],
          [0.187231, 0.414746, 0.556547],
          [0.185556, 0.418570, 0.556753],
          [0.183898, 0.422383, 0.556944],
          [0.182256, 0.426184, 0.557120],
          [0.180629, 0.429975, 0.557282],
          [0.179019, 0.433756, 0.557430],
          [0.177423, 0.437527, 0.557565],
          [0.175841, 0.441290, 0.557685],
          [0.174274, 0.445044, 0.557792],
          [0.172719, 0.448791, 0.557885],
          [0.171176, 0.452530, 0.557965],
          [0.169646, 0.456262, 0.558030],
          [0.168126, 0.459988, 0.558082],
          [0.166617, 0.463708, 0.558119],
          [0.165117, 0.467423, 0.558141],
          [0.163625, 0.471133, 0.558148],
          [0.162142, 0.474838, 0.558140],
          [0.160665, 0.478540, 0.558115],
          [0.159194, 0.482237, 0.558073],
          [0.157729, 0.485932, 0.558013],
          [0.156270, 0.489624, 0.557936],
          [0.154815, 0.493313, 0.557840],
          [0.153364, 0.497000, 0.557724],
          [0.151918, 0.500685, 0.557587],
          [0.150476, 0.504369, 0.557430],
          [0.149039, 0.508051, 0.557250],
          [0.147607, 0.511733, 0.557049],
          [0.146180, 0.515413, 0.556823],
          [0.144759, 0.519093, 0.556572],
          [0.143343, 0.522773, 0.556295],
          [0.141935, 0.526453, 0.555991],
          [0.140536, 0.530132, 0.555659],
          [0.139147, 0.533812, 0.555298],
          [0.137770, 0.537492, 0.554906],
          [0.136408, 0.541173, 0.554483],
          [0.135066, 0.544853, 0.554029],
          [0.133743, 0.548535, 0.553541],
          [0.132444, 0.552216, 0.553018],
          [0.131172, 0.555899, 0.552459],
          [0.129933, 0.559582, 0.551864],
          [0.128729, 0.563265, 0.551229],
          [0.127568, 0.566949, 0.550556],
          [0.126453, 0.570633, 0.549841],
          [0.125394, 0.574318, 0.549086],
          [0.124395, 0.578002, 0.548287],
          [0.123463, 0.581687, 0.547445],
          [0.122606, 0.585371, 0.546557],
          [0.121831, 0.589055, 0.545623],
          [0.121148, 0.592739, 0.544641],
          [0.120565, 0.596422, 0.543611],
          [0.120092, 0.600104, 0.542530],
          [0.119738, 0.603785, 0.541400],
          [0.119512, 0.607464, 0.540218],
          [0.119423, 0.611141, 0.538982],
          [0.119483, 0.614817, 0.537692],
          [0.119699, 0.618490, 0.536347],
          [0.120081, 0.622161, 0.534946],
          [0.120638, 0.625828, 0.533488],
          [0.121380, 0.629492, 0.531973],
          [0.122312, 0.633153, 0.530398],
          [0.123444, 0.636809, 0.528763],
          [0.124780, 0.640461, 0.527068],
          [0.126326, 0.644107, 0.525311],
          [0.128087, 0.647749, 0.523491],
          [0.130067, 0.651384, 0.521608],
          [0.132268, 0.655014, 0.519661],
          [0.134692, 0.658636, 0.517649],
          [0.137339, 0.662252, 0.515571],
          [0.140210, 0.665859, 0.513427],
          [0.143303, 0.669459, 0.511215],
          [0.146616, 0.673050, 0.508936],
          [0.150148, 0.676631, 0.506589],
          [0.153894, 0.680203, 0.504172],
          [0.157851, 0.683765, 0.501686],
          [0.162016, 0.687316, 0.499129],
          [0.166383, 0.690856, 0.496502],
          [0.170948, 0.694384, 0.493803],
          [0.175707, 0.697900, 0.491033],
          [0.180653, 0.701402, 0.488189],
          [0.185783, 0.704891, 0.485273],
          [0.191090, 0.708366, 0.482284],
          [0.196571, 0.711827, 0.479221],
          [0.202219, 0.715272, 0.476084],
          [0.208030, 0.718701, 0.472873],
          [0.214000, 0.722114, 0.469588],
          [0.220124, 0.725509, 0.466226],
          [0.226397, 0.728888, 0.462789],
          [0.232815, 0.732247, 0.459277],
          [0.239374, 0.735588, 0.455688],
          [0.246070, 0.738910, 0.452024],
          [0.252899, 0.742211, 0.448284],
          [0.259857, 0.745492, 0.444467],
          [0.266941, 0.748751, 0.440573],
          [0.274149, 0.751988, 0.436601],
          [0.281477, 0.755203, 0.432552],
          [0.288921, 0.758394, 0.428426],
          [0.296479, 0.761561, 0.424223],
          [0.304148, 0.764704, 0.419943],
          [0.311925, 0.767822, 0.415586],
          [0.319809, 0.770914, 0.411152],
          [0.327796, 0.773980, 0.406640],
          [0.335885, 0.777018, 0.402049],
          [0.344074, 0.780029, 0.397381],
          [0.352360, 0.783011, 0.392636],
          [0.360741, 0.785964, 0.387814],
          [0.369214, 0.788888, 0.382914],
          [0.377779, 0.791781, 0.377939],
          [0.386433, 0.794644, 0.372886],
          [0.395174, 0.797475, 0.367757],
          [0.404001, 0.800275, 0.362552],
          [0.412913, 0.803041, 0.357269],
          [0.421908, 0.805774, 0.351910],
          [0.430983, 0.808473, 0.346476],
          [0.440137, 0.811138, 0.340967],
          [0.449368, 0.813768, 0.335384],
          [0.458674, 0.816363, 0.329727],
          [0.468053, 0.818921, 0.323998],
          [0.477504, 0.821444, 0.318195],
          [0.487026, 0.823929, 0.312321],
          [0.496615, 0.826376, 0.306377],
          [0.506271, 0.828786, 0.300362],
          [0.515992, 0.831158, 0.294279],
          [0.525776, 0.833491, 0.288127],
          [0.535621, 0.835785, 0.281908],
          [0.545524, 0.838039, 0.275626],
          [0.555484, 0.840254, 0.269281],
          [0.565498, 0.842430, 0.262877],
          [0.575563, 0.844566, 0.256415],
          [0.585678, 0.846661, 0.249897],
          [0.595839, 0.848717, 0.243329],
          [0.606045, 0.850733, 0.236712],
          [0.616293, 0.852709, 0.230052],
          [0.626579, 0.854645, 0.223353],
          [0.636902, 0.856542, 0.216620],
          [0.647257, 0.858400, 0.209861],
          [0.657642, 0.860219, 0.203082],
          [0.668054, 0.861999, 0.196293],
          [0.678489, 0.863742, 0.189503],
          [0.688944, 0.865448, 0.182725],
          [0.699415, 0.867117, 0.175971],
          [0.709898, 0.868751, 0.169257],
          [0.720391, 0.870350, 0.162603],
          [0.730889, 0.871916, 0.156029],
          [0.741388, 0.873449, 0.149561],
          [0.751884, 0.874951, 0.143228],
          [0.762373, 0.876424, 0.137064],
          [0.772852, 0.877868, 0.131109],
          [0.783315, 0.879285, 0.125405],
          [0.793760, 0.880678, 0.120005],
          [0.804182, 0.882046, 0.114965],
          [0.814576, 0.883393, 0.110347],
          [0.824940, 0.884720, 0.106217],
          [0.835270, 0.886029, 0.102646],
          [0.845561, 0.887322, 0.099702],
          [0.855810, 0.888601, 0.097452],
          [0.866013, 0.889868, 0.095953],
          [0.876168, 0.891125, 0.095250],
          [0.886271, 0.892374, 0.095374],
          [0.896320, 0.893616, 0.096335],
          [0.906311, 0.894855, 0.098125],
          [0.916242, 0.896091, 0.100717],
          [0.926106, 0.897330, 0.104071],
          [0.935904, 0.898570, 0.108131],
          [0.945636, 0.899815, 0.112838],
          [0.955300, 0.901065, 0.118128],
          [0.964894, 0.902323, 0.123941],
          [0.974417, 0.903590, 0.130215],
          [0.983868, 0.904867, 0.136897],
          [0.993248, 0.906157, 0.143936]]

cvirdis = [[0.000000, 0.135112, 0.304751],
           [0.000000, 0.138068, 0.311105],
           [0.000000, 0.141013, 0.317579],
           [0.000000, 0.143951, 0.323982],
           [0.000000, 0.146877, 0.330479],
           [0.000000, 0.149791, 0.337065],
           [0.000000, 0.152673, 0.343704],
           [0.000000, 0.155377, 0.350500],
           [0.000000, 0.157932, 0.357521],
           [0.000000, 0.160495, 0.364534],
           [0.000000, 0.163058, 0.371608],
           [0.000000, 0.165621, 0.378769],
           [0.000000, 0.168204, 0.385902],
           [0.000000, 0.170800, 0.393100],
           [0.000000, 0.173420, 0.400353],
           [0.000000, 0.176082, 0.407577],
           [0.000000, 0.178802, 0.414764],
           [0.000000, 0.181610, 0.421859],
           [0.000000, 0.184550, 0.428802],
           [0.000000, 0.186915, 0.435532],
           [0.000000, 0.188769, 0.439563],
           [0.000000, 0.190950, 0.441085],
           [0.000000, 0.193366, 0.441561],
           [0.003602, 0.195911, 0.441564],
           [0.017852, 0.198528, 0.441248],
           [0.032110, 0.201199, 0.440785],
           [0.046205, 0.203903, 0.440196],
           [0.058378, 0.206629, 0.439531],
           [0.068968, 0.209372, 0.438863],
           [0.078624, 0.212122, 0.438105],
           [0.087465, 0.214879, 0.437342],
           [0.095645, 0.217643, 0.436593],
           [0.103401, 0.220406, 0.435790],
           [0.110658, 0.223170, 0.435067],
           [0.117612, 0.225935, 0.434308],
           [0.124291, 0.228697, 0.433547],
           [0.130669, 0.231458, 0.432840],
           [0.136830, 0.234216, 0.432148],
           [0.142852, 0.236972, 0.431404],
           [0.148638, 0.239724, 0.430752],
           [0.154261, 0.242475, 0.430120],
           [0.159733, 0.245221, 0.429528],
           [0.165113, 0.247965, 0.428908],
           [0.170362, 0.250707, 0.428325],
           [0.175490, 0.253444, 0.427790],
           [0.180503, 0.256180, 0.427299],
           [0.185453, 0.258914, 0.426788],
           [0.190303, 0.261644, 0.426329],
           [0.195057, 0.264372, 0.425924],
           [0.199764, 0.267099, 0.425497],
           [0.204385, 0.269823, 0.425126],
           [0.208926, 0.272546, 0.424809],
           [0.213431, 0.275266, 0.424480],
           [0.217863, 0.277985, 0.424206],
           [0.222264, 0.280702, 0.423914],
           [0.226598, 0.283419, 0.423678],
           [0.230871, 0.286134, 0.423498],
           [0.235120, 0.288848, 0.423304],
           [0.239312, 0.291562, 0.423167],
           [0.243485, 0.294274, 0.423014],
           [0.247605, 0.296986, 0.422917],
           [0.251675, 0.299698, 0.422873],
           [0.255731, 0.302409, 0.422814],
           [0.259740, 0.305120, 0.422810],
           [0.263738, 0.307831, 0.422789],
           [0.267693, 0.310542, 0.422821],
           [0.271639, 0.313253, 0.422837],
           [0.275513, 0.315965, 0.422979],
           [0.279411, 0.318677, 0.423031],
           [0.283240, 0.321390, 0.423211],
           [0.287065, 0.324103, 0.423373],
           [0.290884, 0.326816, 0.423517],
           [0.294669, 0.329531, 0.423716],
           [0.298421, 0.332247, 0.423973],
           [0.302169, 0.334963, 0.424213],
           [0.305886, 0.337681, 0.424512],
           [0.309601, 0.340399, 0.424790],
           [0.313287, 0.343120, 0.425120],
           [0.316941, 0.345842, 0.425512],
           [0.320595, 0.348565, 0.425889],
           [0.324250, 0.351289, 0.426250],
           [0.327875, 0.354016, 0.426670],
           [0.331474, 0.356744, 0.427144],
           [0.335073, 0.359474, 0.427605],
           [0.338673, 0.362206, 0.428053],
           [0.342246, 0.364939, 0.428559],
           [0.345793, 0.367676, 0.429127],
           [0.349341, 0.370414, 0.429685],
           [0.352892, 0.373153, 0.430226],
           [0.356418, 0.375896, 0.430823],
           [0.359916, 0.378641, 0.431501],
           [0.363446, 0.381388, 0.432075],
           [0.366923, 0.384139, 0.432796],
           [0.370430, 0.386890, 0.433428],
           [0.373884, 0.389646, 0.434209],
           [0.377371, 0.392404, 0.434890],
           [0.380830, 0.395164, 0.435653],
           [0.384268, 0.397928, 0.436475],
           [0.387705, 0.400694, 0.437305],
           [0.391151, 0.403464, 0.438096],
           [0.394568, 0.406236, 0.438986],
           [0.397991, 0.409011, 0.439848],
           [0.401418, 0.411790, 0.440708],
           [0.404820, 0.414572, 0.441642],
           [0.408226, 0.417357, 0.442570],
           [0.411607, 0.420145, 0.443577],
           [0.414992, 0.422937, 0.444578],
           [0.418383, 0.425733, 0.445560],
           [0.421748, 0.428531, 0.446640],
           [0.425120, 0.431334, 0.447692],
           [0.428462, 0.434140, 0.448864],
           [0.431817, 0.436950, 0.449982],
           [0.435168, 0.439763, 0.451134],
           [0.438504, 0.442580, 0.452341],
           [0.441810, 0.445402, 0.453659],
           [0.445148, 0.448226, 0.454885],
           [0.448447, 0.451053, 0.456264],
           [0.451759, 0.453887, 0.457582],
           [0.455072, 0.456718, 0.458976],
           [0.458366, 0.459552, 0.460457],
           [0.461616, 0.462405, 0.461969],
           [0.464947, 0.465241, 0.463395],
           [0.468254, 0.468083, 0.464908],
           [0.471501, 0.470960, 0.466357],
           [0.474812, 0.473832, 0.467681],
           [0.478186, 0.476699, 0.468845],
           [0.481622, 0.479573, 0.469767],
           [0.485141, 0.482451, 0.470384],
           [0.488697, 0.485318, 0.471008],
           [0.492278, 0.488198, 0.471453],
           [0.495913, 0.491076, 0.471751],
           [0.499552, 0.493960, 0.472032],
           [0.503185, 0.496851, 0.472305],
           [0.506866, 0.499743, 0.472432],
           [0.510540, 0.502643, 0.472550],
           [0.514226, 0.505546, 0.472640],
           [0.517920, 0.508454, 0.472707],
           [0.521643, 0.511367, 0.472639],
           [0.525348, 0.514285, 0.472660],
           [0.529086, 0.517207, 0.472543],
           [0.532829, 0.520135, 0.472401],
           [0.536553, 0.523067, 0.472352],
           [0.540307, 0.526005, 0.472163],
           [0.544069, 0.528948, 0.471947],
           [0.547840, 0.531895, 0.471704],
           [0.551612, 0.534849, 0.471439],
           [0.555393, 0.537807, 0.471147],
           [0.559181, 0.540771, 0.470829],
           [0.562972, 0.543741, 0.470488],
           [0.566802, 0.546715, 0.469988],
           [0.570607, 0.549695, 0.469593],
           [0.574417, 0.552682, 0.469172],
           [0.578236, 0.555673, 0.468724],
           [0.582087, 0.558670, 0.468118],
           [0.585916, 0.561674, 0.467618],
           [0.589753, 0.564682, 0.467090],
           [0.593622, 0.567697, 0.466401],
           [0.597469, 0.570718, 0.465821],
           [0.601354, 0.573743, 0.465074],
           [0.605211, 0.576777, 0.464441],
           [0.609105, 0.579816, 0.463638],
           [0.612977, 0.582861, 0.462950],
           [0.616852, 0.585913, 0.462237],
           [0.620765, 0.588970, 0.461351],
           [0.624654, 0.592034, 0.460583],
           [0.628576, 0.595104, 0.459641],
           [0.632506, 0.598180, 0.458668],
           [0.636412, 0.601264, 0.457818],
           [0.640352, 0.604354, 0.456791],
           [0.644270, 0.607450, 0.455886],
           [0.648222, 0.610553, 0.454801],
           [0.652178, 0.613664, 0.453689],
           [0.656114, 0.616780, 0.452702],
           [0.660082, 0.619904, 0.451534],
           [0.664055, 0.623034, 0.450338],
           [0.668008, 0.626171, 0.449270],
           [0.671991, 0.629316, 0.448018],
           [0.675981, 0.632468, 0.446736],
           [0.679979, 0.635626, 0.445424],
           [0.683950, 0.638793, 0.444251],
           [0.687957, 0.641966, 0.442886],
           [0.691971, 0.645145, 0.441491],
           [0.695985, 0.648334, 0.440072],
           [0.700008, 0.651529, 0.438624],
           [0.704037, 0.654731, 0.437147],
           [0.708067, 0.657942, 0.435647],
           [0.712105, 0.661160, 0.434117],
           [0.716177, 0.664384, 0.432386],
           [0.720222, 0.667618, 0.430805],
           [0.724274, 0.670859, 0.429194],
           [0.728334, 0.674107, 0.427554],
           [0.732422, 0.677364, 0.425717],
           [0.736488, 0.680629, 0.424028],
           [0.740589, 0.683900, 0.422131],
           [0.744664, 0.687181, 0.420393],
           [0.748772, 0.690470, 0.418448],
           [0.752886, 0.693766, 0.416472],
           [0.756975, 0.697071, 0.414659],
           [0.761096, 0.700384, 0.412638],
           [0.765223, 0.703705, 0.410587],
           [0.769353, 0.707035, 0.408516],
           [0.773486, 0.710373, 0.406422],
           [0.777651, 0.713719, 0.404112],
           [0.781795, 0.717074, 0.401966],
           [0.785965, 0.720438, 0.399613],
           [0.790116, 0.723810, 0.397423],
           [0.794298, 0.727190, 0.395016],
           [0.798480, 0.730580, 0.392597],
           [0.802667, 0.733978, 0.390153],
           [0.806859, 0.737385, 0.387684],
           [0.811054, 0.740801, 0.385198],
           [0.815274, 0.744226, 0.382504],
           [0.819499, 0.747659, 0.379785],
           [0.823729, 0.751101, 0.377043],
           [0.827959, 0.754553, 0.374292],
           [0.832192, 0.758014, 0.371529],
           [0.836429, 0.761483, 0.368747],
           [0.840693, 0.764962, 0.365746],
           [0.844957, 0.768450, 0.362741],
           [0.849223, 0.771947, 0.359729],
           [0.853515, 0.775454, 0.356500],
           [0.857809, 0.778969, 0.353259],
           [0.862105, 0.782494, 0.350011],
           [0.866421, 0.786028, 0.346571],
           [0.870717, 0.789572, 0.343333],
           [0.875057, 0.793125, 0.339685],
           [0.879378, 0.796687, 0.336241],
           [0.883720, 0.800258, 0.332599],
           [0.888081, 0.803839, 0.328770],
           [0.892440, 0.807430, 0.324968],
           [0.896818, 0.811030, 0.320982],
           [0.901195, 0.814639, 0.317021],
           [0.905589, 0.818257, 0.312889],
           [0.910000, 0.821885, 0.308594],
           [0.914407, 0.825522, 0.304348],
           [0.918828, 0.829168, 0.299960],
           [0.923279, 0.832822, 0.295244],
           [0.927724, 0.836486, 0.290611],
           [0.932180, 0.840159, 0.285880],
           [0.936660, 0.843841, 0.280876],
           [0.941147, 0.847530, 0.275815],
           [0.945654, 0.851228, 0.270532],
           [0.950178, 0.854933, 0.265085],
           [0.954725, 0.858646, 0.259365],
           [0.959284, 0.862365, 0.253563],
           [0.963872, 0.866089, 0.247445],
           [0.968469, 0.869819, 0.241310],
           [0.973114, 0.873550, 0.234677],
           [0.977780, 0.877281, 0.227954],
           [0.982497, 0.881008, 0.220878],
           [0.987293, 0.884718, 0.213336],
           [0.992218, 0.888385, 0.205468],
           [0.994847, 0.892954, 0.203445],
           [0.995249, 0.898384, 0.207561],
           [0.995503, 0.903866, 0.212370],
           [0.995737, 0.909344, 0.217772]]

turbo = [[0.18995, 0.07176, 0.23217],
         [0.19483, 0.08339, 0.26149],
         [0.19956, 0.09498, 0.29024],
         [0.20415, 0.10652, 0.31844],
         [0.20860, 0.11802, 0.34607],
         [0.21291, 0.12947, 0.37314],
         [0.21708, 0.14087, 0.39964],
         [0.22111, 0.15223, 0.42558],
         [0.22500, 0.16354, 0.45096],
         [0.22875, 0.17481, 0.47578],
         [0.23236, 0.18603, 0.50004],
         [0.23582, 0.19720, 0.52373],
         [0.23915, 0.20833, 0.54686],
         [0.24234, 0.21941, 0.56942],
         [0.24539, 0.23044, 0.59142],
         [0.24830, 0.24143, 0.61286],
         [0.25107, 0.25237, 0.63374],
         [0.25369, 0.26327, 0.65406],
         [0.25618, 0.27412, 0.67381],
         [0.25853, 0.28492, 0.69300],
         [0.26074, 0.29568, 0.71162],
         [0.26280, 0.30639, 0.72968],
         [0.26473, 0.31706, 0.74718],
         [0.26652, 0.32768, 0.76412],
         [0.26816, 0.33825, 0.78050],
         [0.26967, 0.34878, 0.79631],
         [0.27103, 0.35926, 0.81156],
         [0.27226, 0.36970, 0.82624],
         [0.27334, 0.38008, 0.84037],
         [0.27429, 0.39043, 0.85393],
         [0.27509, 0.40072, 0.86692],
         [0.27576, 0.41097, 0.87936],
         [0.27628, 0.42118, 0.89123],
         [0.27667, 0.43134, 0.90254],
         [0.27691, 0.44145, 0.91328],
         [0.27701, 0.45152, 0.92347],
         [0.27698, 0.46153, 0.93309],
         [0.27680, 0.47151, 0.94214],
         [0.27648, 0.48144, 0.95064],
         [0.27603, 0.49132, 0.95857],
         [0.27543, 0.50115, 0.96594],
         [0.27469, 0.51094, 0.97275],
         [0.27381, 0.52069, 0.97899],
         [0.27273, 0.53040, 0.98461],
         [0.27106, 0.54015, 0.98930],
         [0.26878, 0.54995, 0.99303],
         [0.26592, 0.55979, 0.99583],
         [0.26252, 0.56967, 0.99773],
         [0.25862, 0.57958, 0.99876],
         [0.25425, 0.58950, 0.99896],
         [0.24946, 0.59943, 0.99835],
         [0.24427, 0.60937, 0.99697],
         [0.23874, 0.61931, 0.99485],
         [0.23288, 0.62923, 0.99202],
         [0.22676, 0.63913, 0.98851],
         [0.22039, 0.64901, 0.98436],
         [0.21382, 0.65886, 0.97959],
         [0.20708, 0.66866, 0.97423],
         [0.20021, 0.67842, 0.96833],
         [0.19326, 0.68812, 0.96190],
         [0.18625, 0.69775, 0.95498],
         [0.17923, 0.70732, 0.94761],
         [0.17223, 0.71680, 0.93981],
         [0.16529, 0.72620, 0.93161],
         [0.15844, 0.73551, 0.92305],
         [0.15173, 0.74472, 0.91416],
         [0.14519, 0.75381, 0.90496],
         [0.13886, 0.76279, 0.89550],
         [0.13278, 0.77165, 0.88580],
         [0.12698, 0.78037, 0.87590],
         [0.12151, 0.78896, 0.86581],
         [0.11639, 0.79740, 0.85559],
         [0.11167, 0.80569, 0.84525],
         [0.10738, 0.81381, 0.83484],
         [0.10357, 0.82177, 0.82437],
         [0.10026, 0.82955, 0.81389],
         [0.09750, 0.83714, 0.80342],
         [0.09532, 0.84455, 0.79299],
         [0.09377, 0.85175, 0.78264],
         [0.09287, 0.85875, 0.77240],
         [0.09267, 0.86554, 0.76230],
         [0.09320, 0.87211, 0.75237],
         [0.09451, 0.87844, 0.74265],
         [0.09662, 0.88454, 0.73316],
         [0.09958, 0.89040, 0.72393],
         [0.10342, 0.89600, 0.71500],
         [0.10815, 0.90142, 0.70599],
         [0.11374, 0.90673, 0.69651],
         [0.12014, 0.91193, 0.68660],
         [0.12733, 0.91701, 0.67627],
         [0.13526, 0.92197, 0.66556],
         [0.14391, 0.92680, 0.65448],
         [0.15323, 0.93151, 0.64308],
         [0.16319, 0.93609, 0.63137],
         [0.17377, 0.94053, 0.61938],
         [0.18491, 0.94484, 0.60713],
         [0.19659, 0.94901, 0.59466],
         [0.20877, 0.95304, 0.58199],
         [0.22142, 0.95692, 0.56914],
         [0.23449, 0.96065, 0.55614],
         [0.24797, 0.96423, 0.54303],
         [0.26180, 0.96765, 0.52981],
         [0.27597, 0.97092, 0.51653],
         [0.29042, 0.97403, 0.50321],
         [0.30513, 0.97697, 0.48987],
         [0.32006, 0.97974, 0.47654],
         [0.33517, 0.98234, 0.46325],
         [0.35043, 0.98477, 0.45002],
         [0.36581, 0.98702, 0.43688],
         [0.38127, 0.98909, 0.42386],
         [0.39678, 0.99098, 0.41098],
         [0.41229, 0.99268, 0.39826],
         [0.42778, 0.99419, 0.38575],
         [0.44321, 0.99551, 0.37345],
         [0.45854, 0.99663, 0.36140],
         [0.47375, 0.99755, 0.34963],
         [0.48879, 0.99828, 0.33816],
         [0.50362, 0.99879, 0.32701],
         [0.51822, 0.99910, 0.31622],
         [0.53255, 0.99919, 0.30581],
         [0.54658, 0.99907, 0.29581],
         [0.56026, 0.99873, 0.28623],
         [0.57357, 0.99817, 0.27712],
         [0.58646, 0.99739, 0.26849],
         [0.59891, 0.99638, 0.26038],
         [0.61088, 0.99514, 0.25280],
         [0.62233, 0.99366, 0.24579],
         [0.63323, 0.99195, 0.23937],
         [0.64362, 0.98999, 0.23356],
         [0.65394, 0.98775, 0.22835],
         [0.66428, 0.98524, 0.22370],
         [0.67462, 0.98246, 0.21960],
         [0.68494, 0.97941, 0.21602],
         [0.69525, 0.97610, 0.21294],
         [0.70553, 0.97255, 0.21032],
         [0.71577, 0.96875, 0.20815],
         [0.72596, 0.96470, 0.20640],
         [0.73610, 0.96043, 0.20504],
         [0.74617, 0.95593, 0.20406],
         [0.75617, 0.95121, 0.20343],
         [0.76608, 0.94627, 0.20311],
         [0.77591, 0.94113, 0.20310],
         [0.78563, 0.93579, 0.20336],
         [0.79524, 0.93025, 0.20386],
         [0.80473, 0.92452, 0.20459],
         [0.81410, 0.91861, 0.20552],
         [0.82333, 0.91253, 0.20663],
         [0.83241, 0.90627, 0.20788],
         [0.84133, 0.89986, 0.20926],
         [0.85010, 0.89328, 0.21074],
         [0.85868, 0.88655, 0.21230],
         [0.86709, 0.87968, 0.21391],
         [0.87530, 0.87267, 0.21555],
         [0.88331, 0.86553, 0.21719],
         [0.89112, 0.85826, 0.21880],
         [0.89870, 0.85087, 0.22038],
         [0.90605, 0.84337, 0.22188],
         [0.91317, 0.83576, 0.22328],
         [0.92004, 0.82806, 0.22456],
         [0.92666, 0.82025, 0.22570],
         [0.93301, 0.81236, 0.22667],
         [0.93909, 0.80439, 0.22744],
         [0.94489, 0.79634, 0.22800],
         [0.95039, 0.78823, 0.22831],
         [0.95560, 0.78005, 0.22836],
         [0.96049, 0.77181, 0.22811],
         [0.96507, 0.76352, 0.22754],
         [0.96931, 0.75519, 0.22663],
         [0.97323, 0.74682, 0.22536],
         [0.97679, 0.73842, 0.22369],
         [0.98000, 0.73000, 0.22161],
         [0.98289, 0.72140, 0.21918],
         [0.98549, 0.71250, 0.21650],
         [0.98781, 0.70330, 0.21358],
         [0.98986, 0.69382, 0.21043],
         [0.99163, 0.68408, 0.20706],
         [0.99314, 0.67408, 0.20348],
         [0.99438, 0.66386, 0.19971],
         [0.99535, 0.65341, 0.19577],
         [0.99607, 0.64277, 0.19165],
         [0.99654, 0.63193, 0.18738],
         [0.99675, 0.62093, 0.18297],
         [0.99672, 0.60977, 0.17842],
         [0.99644, 0.59846, 0.17376],
         [0.99593, 0.58703, 0.16899],
         [0.99517, 0.57549, 0.16412],
         [0.99419, 0.56386, 0.15918],
         [0.99297, 0.55214, 0.15417],
         [0.99153, 0.54036, 0.14910],
         [0.98987, 0.52854, 0.14398],
         [0.98799, 0.51667, 0.13883],
         [0.98590, 0.50479, 0.13367],
         [0.98360, 0.49291, 0.12849],
         [0.98108, 0.48104, 0.12332],
         [0.97837, 0.46920, 0.11817],
         [0.97545, 0.45740, 0.11305],
         [0.97234, 0.44565, 0.10797],
         [0.96904, 0.43399, 0.10294],
         [0.96555, 0.42241, 0.09798],
         [0.96187, 0.41093, 0.09310],
         [0.95801, 0.39958, 0.08831],
         [0.95398, 0.38836, 0.08362],
         [0.94977, 0.37729, 0.07905],
         [0.94538, 0.36638, 0.07461],
         [0.94084, 0.35566, 0.07031],
         [0.93612, 0.34513, 0.06616],
         [0.93125, 0.33482, 0.06218],
         [0.92623, 0.32473, 0.05837],
         [0.92105, 0.31489, 0.05475],
         [0.91572, 0.30530, 0.05134],
         [0.91024, 0.29599, 0.04814],
         [0.90463, 0.28696, 0.04516],
         [0.89888, 0.27824, 0.04243],
         [0.89298, 0.26981, 0.03993],
         [0.88691, 0.26152, 0.03753],
         [0.88066, 0.25334, 0.03521],
         [0.87422, 0.24526, 0.03297],
         [0.86760, 0.23730, 0.03082],
         [0.86079, 0.22945, 0.02875],
         [0.85380, 0.22170, 0.02677],
         [0.84662, 0.21407, 0.02487],
         [0.83926, 0.20654, 0.02305],
         [0.83172, 0.19912, 0.02131],
         [0.82399, 0.19182, 0.01966],
         [0.81608, 0.18462, 0.01809],
         [0.80799, 0.17753, 0.01660],
         [0.79971, 0.17055, 0.01520],
         [0.79125, 0.16368, 0.01387],
         [0.78260, 0.15693, 0.01264],
         [0.77377, 0.15028, 0.01148],
         [0.76476, 0.14374, 0.01041],
         [0.75556, 0.13731, 0.00942],
         [0.74617, 0.13098, 0.00851],
         [0.73661, 0.12477, 0.00769],
         [0.72686, 0.11867, 0.00695],
         [0.71692, 0.11268, 0.00629],
         [0.70680, 0.10680, 0.00571],
         [0.69650, 0.10102, 0.00522],
         [0.68602, 0.09536, 0.00481],
         [0.67535, 0.08980, 0.00449],
         [0.66449, 0.08436, 0.00424],
         [0.65345, 0.07902, 0.00408],
         [0.64223, 0.07380, 0.00401],
         [0.63082, 0.06868, 0.00401],
         [0.61923, 0.06367, 0.00410],
         [0.60746, 0.05878, 0.00427],
         [0.59550, 0.05399, 0.00453],
         [0.58336, 0.04931, 0.00486],
         [0.57103, 0.04474, 0.00529],
         [0.55852, 0.04028, 0.00579],
         [0.54583, 0.03593, 0.00638],
         [0.53295, 0.03169, 0.00705],
         [0.51989, 0.02756, 0.00780],
         [0.50664, 0.02354, 0.00863],
         [0.49321, 0.01963, 0.00955],
         [0.47960, 0.01583, 0.01055]]

# from matplotlib import cm

cm_gpu_store = dotdict()
cm_cpu_store = dotdict()
cm_cpu_store.linear = None
cm_cpu_store.jet = jet
cm_cpu_store.turbo = turbo
cm_cpu_store.virdis = virdis
cm_cpu_store.cvirdis = cvirdis


def colormap(v: torch.Tensor, cm: str = 'virdis'):
    """
    Create colormap from linear mapping segments

    segmentdata argument is a dictionary with a red, green and blue
    entries. Each entry should be a list of *x*, *y0*, *y1* tuples,
    forming rows in a table. Entries for alpha are optional.

    Example: suppose you want red to increase from 0 to 1 over
    the bottom half, green to do the same over the middle half,
    and blue over the top half.  Then you would use::

        jet = {'red':   ((0.00, 0, 0),
                         (0.35, 0, 0),
                         (0.66, 1, 1),
                         (0.89, 1, 1),
                         (1.00, 0.5, 0.5)),
               'green': ((0.000, 0, 0),
                         (0.125, 0, 0),
                         (0.375, 1, 1),
                         (0.640, 1, 1),
                         (0.910, 0, 0),
                         (1.000, 0, 0)),
               'blue':  ((0.00, 0.5, 0.5),
                         (0.11, 1, 1),
                         (0.34, 1, 1),
                         (0.65, 0, 0),
                         (1.00, 0, 0))}

    Each row in the table for a given color is a sequence of
    *x*, *y0*, *y1* tuples.  In each sequence, *x* must increase
    monotonically from 0 to 1.  For any input value *z* falling
    between *x[i]* and *x[i+1]*, the output value of a given color
    will be linearly interpolated between *y1[i]* and *y0[i+1]*::

        row i:   x  y0  y1
                       /
                      /
        row i+1: x  y0  y1

    Hence y0 in the first row and y1 in the last row are never used.

    See Also
    --------
    LinearSegmentedColormap.from_list
        Static method; factory function for generating a smoothly-varying
        LinearSegmentedColormap.
    """

    eps = 1e-6  # consider float32
    v = v.clip(0 + eps, 1 - eps)  # B, P
    if v.shape[-1] != 1:
        v = v[..., None]  # add last dimension
    if cm not in cm_gpu_store:
        cm_gpu_store[cm] = dotdict()
    dv = str(v.device)
    if dv not in cm_gpu_store[cm]:
        mapping = cm_cpu_store[cm]
        if isinstance(mapping, dict):
            # mapping = {k: [torch.as_tensor(v, device=dv) for v in l] for k, l in mapping.items()}
            mapping = torch.stack([torch.as_tensor([mapping['red'][i],
                                                    mapping['green'][i],
                                                    mapping['blue'][i]], device=dv)
                                   for i in range(len(mapping['red']))])
            function = colormap_dict
        elif isinstance(mapping, list):
            mapping = torch.stack([torch.as_tensor(v, device=dv) for v in mapping])  # N, 3
            function = colormap_list
        else:
            mapping = mapping
            function = colormap_linear
        cm_gpu_store[cm][dv] = mapping, function
    cm, fn = cm_gpu_store[cm][dv]

    return fn(v, cm)


@torch.jit.script
def colormap_linear(v: torch.Tensor, cm: NoneType = None):
    return v.expand(v.shape[:-1] + (3,))


@torch.jit.script  # TODO: Check if this will help
def colormap_dict(v: torch.Tensor, cm: torch.Tensor):
    # number, rgb, value+start+end
    color = v.new_zeros(v.shape[:-1] + (3,))  # B, P, 3
    for i in range(len(cm) - 1):
        b0 = cm[i, :, 0]
        b1 = cm[i + 1, :, 0]
        c0 = cm[i, :, 2]
        c1 = cm[i + 1, :, 1]
        interp = (v - b0) / (b1 - b0)  # interpolated value
        valid = (interp >= 0.0) & (interp <= 1.0)  # zero or one
        color += (c1 * interp + c0 * (1 - interp)) * valid
    return color


@torch.jit.script  # TODO: Check if this will help
def colormap_list(v: torch.Tensor, cm: torch.Tensor):
    color = v.new_zeros(v.shape[:-1] + (3,))  # B, P, 3
    for i in range(len(cm) - 1):
        b0 = i / (len(cm) - 1)
        b1 = (i + 1) / (len(cm) - 1)
        c0 = cm[i]
        c1 = cm[i + 1]
        interp = (v - b0) / (b1 - b0)  # interpolated value
        valid = (interp >= 0.0) & (interp <= 1.0)  # zero or one
        color += (c1 * interp + c0 * (1 - interp)) * valid
    return color


def yuv_to_rgb(x):
    y, u, v = x.split([1, 1, 1], dim=-1)
    r: torch.Tensor = y + 1.14 * v  # coefficient for g is 0
    g: torch.Tensor = y + -0.396 * u - 0.581 * v
    b: torch.Tensor = y + 2.029 * u  # coefficient for b is 0
    rgb = torch.cat([r, g, b], dim=-1)
    return rgb


def rgb_to_yuv(x):
    r, g, b = x.split([1, 1, 1], dim=-1)

    y: torch.Tensor = 0.299 * r + 0.587 * g + 0.114 * b
    u: torch.Tensor = -0.147 * r - 0.289 * g + 0.436 * b
    v: torch.Tensor = 0.615 * r - 0.515 * g - 0.100 * b

    yuv = torch.cat([y, u, v], dim=-1)
    return yuv


def image_derivative(img: torch.Tensor, mode='sobel', normalized=True) -> torch.Tensor:
    import kornia
    # compute image space derivative, x right, y down (though, in sampling y should be up)
    # normalized to uv space, thus diff in uv * idr approximates new color value
    # B, H, W, 3
    img_bchw = img.permute(0, 3, 1, 2)
    img_bchw = img_bchw.flip(2)  # flip y dim
    idr_bc2hw = kornia.filters.spatial_gradient(img_bchw, mode, normalized=normalized)
    idr_bc2hw = idr_bc2hw.flip(3)  # flip back y dim
    idr = idr_bc2hw.permute(0, 3, 4, 1, 2)

    if normalized:
        idr[..., 0] = idr[..., 0] * img.shape[2]  # x -> u
        idr[..., 1] = idr[..., 1] * img.shape[1]  # y -> v

    return idr


def image_pyramid(input: torch.Tensor, max_level: int = 4) -> List[torch.Tensor]:
    import kornia
    pyramid = kornia.geometry.transform.build_pyramid(input.permute(0, 3, 1, 2), max_level=max_level, align_corners=False)
    pyramid = [v.permute(0, 2, 3, 1) for v in pyramid]
    return pyramid


def variance_of_laplacian(img: torch.Tensor):
    import kornia
    sh = img.shape
    if len(sh) == 4 and (sh[-1] == 3 or sh[-1] == 1):  # assumption
        img = img.permute(0, 3, 1, 2)
    elif len(sh) == 3 and (sh[-1] == 3 or sh[-1] == 1):  # assumption
        img = img.permute(2, 0, 1)[None]
    laplacian = kornia.filters.laplacian(img)  # B, C, H, W
    return laplacian.var()
